
#include "stdafx.h"

#include <cstdio>
#include <string>
#include <sstream>
#include <iostream>
#include<set>
#include<map>

#include "clang/AST/ASTConsumer.h"
#include "clang/AST/RecursiveASTVisitor.h"
#include "clang/Basic/Diagnostic.h"
#include "clang/Basic/FileManager.h"
#include "clang/Basic/SourceManager.h"
#include "clang/Basic/TargetOptions.h"
#include "clang/Basic/TargetInfo.h"
#include "clang/Frontend/CompilerInstance.h"
#include "clang/Lex/DirectoryLookup.h"
#include "clang/Lex/HeaderSearch.h"
#include "clang/Lex/Preprocessor.h"
#include "clang/Parse/ParseAST.h"
#include "clang/Rewrite/Core/Rewriter.h"
#include "clang/Rewrite/Frontend/Rewriters.h"
#include "llvm/Support/Host.h"
#include "llvm/Support/raw_ostream.h"

#include "llvm/Support/Commandline.h"
#include "clang/Frontend/FrontendActions.h"
#include "clang/Tooling/CompilationDatabase.h"
#include "clang/Tooling/Tooling.h"

using namespace clang;
using namespace std;
using namespace clang::tooling;
using namespace llvm;

cl::opt<std::string> BuildPath(
	"p",
	cl::desc("<build-path>"),
	cl::Optional);

cl::list<std::string> SourcePaths(
	cl::Positional,
	cl::desc("<source0> [... <sourceN>]"),
	cl::OneOrMore);

enum Mode {INSERT, REPLACE};

class MyRewriter: public Rewriter
{
public:
	void setMode(Mode m) {
		mode =  m;
	}
	Mode getMode() {
		return mode;
	}
	bool isNameLookedFor(string s) {
		return (TheNamesLookedFor.find(s) != TheNamesLookedFor.end());
	}
	bool isNameReplaced(string s) {
		return (TheNamesReplaced.find(s) != TheNamesReplaced.end());
	}

	string getReplacementFor(string s) {
		return TheNamesReplaced[s];
	}

	void AddNameLookedFor(string s) {
		TheNamesLookedFor.insert(s);
	}

	void AddNameToReplace(string s, string replacement) {
		TheNamesReplaced[s] = replacement;
	}

private:
	Mode mode;
	set<string> TheNamesLookedFor;
	map<string, string> TheNamesReplaced;
};


// By implementing RecursiveASTVisitor, we can specify which AST nodes
// we're interested in by overriding relevant methods.
class MyASTVisitor : public RecursiveASTVisitor<MyASTVisitor>
{
public:
    MyASTVisitor(MyRewriter &R): TheRewriter(R)
    {}

    bool VisitFunctionDecl(FunctionDecl *f) {

		SourceRange SR = f->getSourceRange();

		QualType QT = f->getResultType();
        string TypeStr = QT.getAsString();

		string FuncName = f ->getNameAsString();

		if (TheRewriter.getMode() == INSERT) {

			SourceLocation SL = SR.getBegin();
			SStream <<  "/* Declaration of function " << FuncName << " returning "
                     << TypeStr << "*/\n";
			if (TheRewriter.isNameLookedFor(FuncName)) {
				TheRewriter.InsertText(SL, SStream.str(),  true, true);
			}			
		}

		if (TheRewriter.getMode() == REPLACE) {

			SourceLocation SL = f->getLocation();
			int offset = FuncName.length() - 1;
			SourceRange SRN = SourceRange(SL, SL.getLocWithOffset(offset));
			if (TheRewriter.isNameReplaced(FuncName)) {
				TheRewriter.ReplaceText(SRN, TheRewriter.getReplacementFor(FuncName));
			}
		};
		clearSStream();
		return true;
    }

	bool VisitVarDecl(VarDecl *v) {

		SourceRange SR = v->getSourceRange();

        string TheVarName = v->getNameAsString();
		QualType QT = v ->getType();
		string TheVarType = QT.getAsString();

		Expr * InitExpr = (v->getInit());
		string  InitExprType;
		if (InitExpr) {
			InitExpr =  InitExpr->IgnoreImpCasts();
			QualType QT = InitExpr->getType();
			InitExprType = QT.getAsString();
		}

		if ((TheRewriter.getMode() == INSERT) && (TheRewriter.isNameLookedFor(TheVarName))) {

			SourceLocation SL = SR.getBegin();
			SStream << "/*Var " << TheVarName <<" of type " << TheVarType << ". ";
			if (InitExpr) {
				SStream << "Initialized with: " << InitExprType;
			}
			SStream << "*/";
			TheRewriter.InsertText(SL, SStream.str(), true, true);			
		}

		if ((TheRewriter.getMode() == REPLACE) && (TheRewriter.isNameReplaced(TheVarName))) {

			SourceLocation SL = v->getLocation();
			int offset = TheVarName.length() - 1;
			SourceRange SR = SourceRange(SL, SL.getLocWithOffset(offset));
			TheRewriter.ReplaceText(SR, TheRewriter.getReplacementFor(TheVarName));
		};

		clearSStream();
		return true;
    }



	bool VisitFieldDecl(FieldDecl *TheFieldDecl) {

		SourceRange SR = TheFieldDecl->getSourceRange();
		string FieldName = TheFieldDecl->getNameAsString();
		string ParentName = TheFieldDecl->getParent()->getNameAsString();
		if (TheRewriter.getMode() == INSERT) {

			SourceLocation SL = SR.getBegin();
			SStream <<  "/* Declaration of field " << FieldName << " of record " << ParentName << "*/\n";
			if (TheRewriter.isNameLookedFor(FieldName)) {
				TheRewriter.InsertText(SL, SStream.str(),  true, true);
			}			
		}

		if (TheRewriter.getMode() == REPLACE) {

			SourceLocation SL = TheFieldDecl->getLocation();
			int offset = FieldName.length() - 1;
			SourceRange SRN = SourceRange(SL, SL.getLocWithOffset(offset));
			if (TheRewriter.isNameReplaced(FieldName)) {
				TheRewriter.ReplaceText(SRN, TheRewriter.getReplacementFor(FieldName));
			}
		};
		clearSStream();
		return true;
    }


	bool VisitBinaryOperator(BinaryOperator *binop) {

		SourceLocation ST = binop-> getSourceRange().getBegin();
        Expr* lhs = binop->getLHS();
		Expr* rhs = (binop->getRHS())->IgnoreImpCasts();
		QualType QT = lhs->getType();
		string  lhsType = QT.getAsString();
		QT = rhs->getType();
		string  rhsType = QT.getAsString();

		clearSStream();
		return true;
	}

	
	bool VisitDeclRefExpr(DeclRefExpr * TheDeclRefExpr) {
		ValueDecl* TheDecl  = TheDeclRefExpr->getDecl();
		SourceRange SR = TheDeclRefExpr->getSourceRange();
		if ((TheDecl->isFunctionOrFunctionTemplate()))
		{
			string FuncName = TheDecl->getNameAsString();
			if ((TheRewriter.getMode() == INSERT) && (TheRewriter.isNameLookedFor(FuncName))) {
				
				SourceLocation SL = SR.getBegin();
				SStream << "/*Called function " << FuncName  << "*/";
				TheRewriter.InsertText(SL, SStream.str(), true, true);		
			}

			if ((TheRewriter.getMode() == REPLACE) && (TheRewriter.isNameReplaced(FuncName))) {
				SourceLocation SL = TheDeclRefExpr->getLocation();
				int offset = FuncName.length() - 1;
				SourceRange SRN = SourceRange(SL, SL.getLocWithOffset(offset));			
				
				TheRewriter.ReplaceText(SRN, TheRewriter.getReplacementFor(FuncName));
			}
		}
		else
		{
			
			string DeclName = TheDecl->getNameAsString();
			if ((TheRewriter.getMode() == INSERT) && (TheRewriter.isNameLookedFor(DeclName))) {

					SourceLocation SL = SR.getBegin();						
					SStream << "/*Var " << DeclName  << " used*/";
					TheRewriter.InsertText(SL, SStream.str(), true, true);	
				
			}

			if ((TheRewriter.getMode() == REPLACE) && (TheRewriter.isNameReplaced(DeclName))) {
				SourceLocation SL = TheDeclRefExpr->getLocStart();
				int offset = DeclName.length() - 1;
				SourceRange SRN = SourceRange(SL, SL.getLocWithOffset(offset));
				TheRewriter.ReplaceText(SRN, TheRewriter.getReplacementFor(DeclName));
			}
			clearSStream();
		}
		return true;
	};

	
	bool VisitCallExpr(CallExpr * TheCallExpr) {
		
		Decl* TheDecl  = TheCallExpr->getCalleeDecl();		
		SourceRange SR = TheCallExpr->getSourceRange();		
		
		if (TheDecl->isFunctionOrFunctionTemplate()) {
		
			FunctionDecl *TheFuncDecl = static_cast<FunctionDecl *>(TheDecl);
			
			DeclarationName DeclName = TheFuncDecl->getNameInfo().getName();
            string FuncName = DeclName.getAsString();	
			
			/*
            if ((TheRewriter.getMode() == INSERT) && (TheRewriter.isNameLookedFor(FuncName))) {
				
				SourceLocation SL = SR.getBegin();
				SStream << "/*Called function " << FuncName  << "*//*";
				TheRewriter.InsertText(SL, SStream.str(), true, true);		
			}

			if ((TheRewriter.getMode() == REPLACE) && (TheRewriter.isNameReplaced(FuncName))) {
				SourceLocation SL = TheCallExpr->getExprLoc();
				int offset = FuncName.length() - 1;
				SourceRange SRN = SourceRange(SL, SL.getLocWithOffset(offset));
				/T/heRewriter.InsertText(SL, "/*here", true, true);
				
				TheRewriter.ReplaceText(SRN, TheRewriter.getReplacementFor(FuncName));
			} */
		}		
		
		Expr ** TheArgs = TheCallExpr->getArgs();
		for (CallExpr::arg_iterator i = TheCallExpr->arg_begin(); i != TheCallExpr->arg_end(); ++i) {

			Expr * CurrArg = i->IgnoreImpCasts();			
			QualType QT = CurrArg->getType();			
			string  CurrArgType = QT.getAsString();			
			
		}
		
		return true;
	};

	
	bool VisitMemberExpr(MemberExpr * TheMemberExpr) {
		
		ValueDecl* TheMemberDecl  = TheMemberExpr->getMemberDecl();		
		SourceRange SR = TheMemberExpr->getSourceRange();			
			
		string MemberName = TheMemberDecl->getNameAsString(); 
		
		
			if ((TheRewriter.getMode() == INSERT) && (TheRewriter.isNameLookedFor(MemberName))) {
				
				SourceLocation SL = SR.getBegin();
				if (TheMemberDecl->isFunctionOrFunctionTemplate())
				SStream << "/*Called method " << MemberName  << "*/";
				else
				{
					SStream << "/*Field " << MemberName  <<  "*/";
				}

				TheRewriter.InsertText(SL, SStream.str(), true, true);		
			}

			if ((TheRewriter.getMode() == REPLACE) && (TheRewriter.isNameReplaced(MemberName))) {
				SourceLocation SL = TheMemberExpr->getMemberLoc();
				int offset = MemberName.length() - 1;
				SourceRange SRN = SourceRange(SL, SL.getLocWithOffset(offset));
				//TheRewriter.InsertText(SL, "/*here", true, true);				
				TheRewriter.ReplaceText(SRN, TheRewriter.getReplacementFor(MemberName));
			} 			
			
		clearSStream();		
		return true;
	};
	

	bool VisitCXXRecordDecl(CXXRecordDecl *r) {

		string RecordName = r->getNameAsString();
		if ((TheRewriter.getMode() == INSERT) && (TheRewriter.isNameLookedFor(RecordName) )) {

			SourceLocation SL = r->getSourceRange().getBegin();
			if (r->isStruct()) {

				SStream << "//Struct " << RecordName << " found\n";
				TheRewriter.InsertText(SL, SStream.str(), true, true);
			}
			if (r->isClass()) {
				SStream << "//Class " << RecordName << " found\n";
				TheRewriter.InsertText(SL, SStream.str(), true, true);
			}
		}
		if ((TheRewriter.getMode() == REPLACE) && (TheRewriter.isNameReplaced(RecordName))) {

			SourceLocation SL = r->getLocation();
			int offset = RecordName.length() - 1;
			SourceRange SR = SourceRange(SL, SL.getLocWithOffset(offset));
			TheRewriter.ReplaceText(SR, TheRewriter.getReplacementFor(RecordName));
		};

		clearSStream();
        return true;
    }

private:
    void AddBraces(Stmt *s);
	MyRewriter &TheRewriter;
	stringstream SStream;
	void clearSStream() {
		SStream.str("");
		SStream.clear();

	}
};

// Implementation of the ASTConsumer interface for reading an AST produced
// by the Clang parser.
class MyASTConsumer : public ASTConsumer
{
public:
    MyASTConsumer(MyRewriter &R)
        : Visitor(R)

	{}

    // Override the method that gets called for each parsed top-level
    // declaration.
    virtual bool HandleTopLevelDecl(DeclGroupRef DR) {
        for (DeclGroupRef::iterator b = DR.begin(), e = DR.end();
             b != e; ++b)
            // Traverse the declaration using our AST visitor.
            Visitor.TraverseDecl(*b);
        return true;
    }

private:
    MyASTVisitor Visitor;

};

int runBatch_main(int argc, const char **argv)
{
	llvm::OwningPtr<CompilationDatabase> Compilations(
		FixedCompilationDatabase::loadFromCommandLine(argc, argv));
	cl::ParseCommandLineOptions(argc,argv);
	if(!Compilations){
		std::string ErrorMessage;
		if(!BuildPath.empty()){
			Compilations.reset(
				CompilationDatabase::autoDetectFromDirectory(BuildPath, ErrorMessage));
		}else{
			Compilations.reset(CompilationDatabase::autoDetectFromSource(
				SourcePaths[0], ErrorMessage));
		}
		if (!Compilations)
			llvm::report_fatal_error(ErrorMessage);
	}
	ClangTool Tool(*Compilations, SourcePaths);
	return Tool.run(newFrontendActionFactory<clang::SyntaxOnlyAction>());
}


int main()
{
    // CompilerInstance will hold the instance of the Clang compiler for us,
    // managing the various objects needed to run the compiler.
    CompilerInstance TheCompInst;
	TheCompInst.getLangOpts().CPlusPlus = 1;
    TheCompInst.createDiagnostics(0, 0);

    // Initialize target info with the default triple for our platform.
    TargetOptions TO;
    TO.Triple = llvm::sys::getDefaultTargetTriple();
    TargetInfo *TI = TargetInfo::CreateTargetInfo(
        TheCompInst.getDiagnostics(), TO);
    TheCompInst.setTarget(TI);

    TheCompInst.createFileManager();
    FileManager &FileMgr = TheCompInst.getFileManager();
    TheCompInst.createSourceManager(FileMgr);
    SourceManager &SourceMgr = TheCompInst.getSourceManager();

    TheCompInst.createPreprocessor();
    TheCompInst.createASTContext();

   	MyRewriter TheRewriter;
    TheRewriter.setSourceMgr(SourceMgr, TheCompInst.getLangOpts());

	HeaderSearch & HS = TheCompInst.getPreprocessor().getHeaderSearchInfo();
    TheCompInst.getDiagnosticClient().BeginSourceFile(
        TheCompInst.getLangOpts(),
        &TheCompInst.getPreprocessor());

	//HeaderSearch HSearch(FileMgr, TheCompInst.getDiagnostics(), TheCompInst.getLangOpts(),TI);

	/*vector<DirectoryLookup> v;
	const DirectoryEntry * d[10];
	d[0]= FileMgr.getDirectory("../include");
	v.push_back(DirectoryLookup(d[0], SrcMgr::C_User, true, false));
	d[1] = FileMgr.getDirectory("miranda-read-only/miranda/src/core");
	v.push_back(DirectoryLookup(d[1], SrcMgr::C_User, true, false));
	d[2] = FileMgr.getDirectory("miranda-read-only/miranda/src/modules");
	v.push_back(DirectoryLookup(d[2], SrcMgr::C_User,true,false));
	HS.SetSearchPaths(v,3,3,false);
	//HSearch.AddSearchPath(DirectoryLookup(d[1], SrcMgr::C_User,true,false),false);
	*/

	// Set the main file handled by the source manager to the input file.
	string TheFileName;
	cout << "Enter file name\n";
	getline(cin,TheFileName);
    const FileEntry *FileIn = FileMgr.getFile(TheFileName);
    SourceMgr.createMainFileID(FileIn);


	cout << "Would you like to insert comments (print letter i) or to replace some text (print r)?\n";
	string ModeChoosed ;	
	getline(cin, ModeChoosed);
	istringstream s (ModeChoosed);
	s >> ModeChoosed ;
	if (ModeChoosed == "i") TheRewriter.setMode(INSERT);
	if (ModeChoosed == "r") TheRewriter.setMode(REPLACE);
	cin.clear();
    // Create an AST consumer instance which is going to get called by
    // ParseAST.
    if (TheRewriter.getMode() == INSERT) {
		cout << "What names (of functions or variables) should I look for?\n";
		string s;
		getline(cin, s);
		istringstream iss(s);
		string sub;
		while (iss >> sub) {
			TheRewriter.AddNameLookedFor(sub);
		}
	}

	 if (TheRewriter.getMode() == REPLACE) {
		cout << "Print name (of a function or a variable) you want to replace and  (after a space) the replacement\n";
		string s;
		string curr_name, replacement;
		getline(cin, s);
		istringstream iss(s);
		iss >> curr_name;
		iss >> replacement;
		TheRewriter.AddNameToReplace(curr_name, replacement);
	}

	MyASTConsumer TheConsumer(TheRewriter);
    ParseAST(TheCompInst.getPreprocessor(), &TheConsumer,
             TheCompInst.getASTContext());

    const RewriteBuffer *RewriteBuf =
        TheRewriter.getRewriteBufferFor(SourceMgr.getMainFileID());
	TheRewriter.overwriteChangedFiles();

    return 0;
}
